#include <stdio.h>
#include <syslog.h>
#include <errno.h>
#include <string.h>
#include <netinet/in.h>
#include <unistd.h>
#include <arpa/inet.h>
#include <linux/netfilter.h>
#include <libnetfilter_queue/libnetfilter_queue.h>
#include <linux/ip.h>
#include <linux/tcp.h>

unsigned int		g_all_count		= 1;
unsigned int		g_index			= 0;
unsigned int		g_port			= 80;

int cb(struct nfq_q_handle *qh, struct nfgenmsg *nfmsg,
		struct nfq_data *nfa, void *data);
void nfq_listen(void);
void showhelp(void);

/*
 * 入口函数
 */
int main(int argc, char* argv[])
{
	// 解析参数
	extern char* optarg;
	extern int optind, opterr, optopt;
	char			ch;

	opterr = 0;
	while( (ch = getopt(argc, argv, "n:p:")) != -1 )
	{
		switch(ch)
		{
		case 'n':
			g_all_count = atoi(optarg);
			break;
		case 'p':
			g_port = htons( atoi(optarg) );
			break;
		case '?':
			showhelp();	
			return 0;
		}
	}
	// 开循环
	openlog("ipbalance", LOG_PID, LOG_USER);
	syslog(LOG_INFO, "begin listen");
	nfq_listen();
	syslog(LOG_INFO, "end listen");
	closelog();
	return 0;
}

/*
 * 显示帮助
 */
void showhelp(void)
{
	printf("ipbalance [option]\r\n"
			"\t-n [number]: set all count(defailt 0)\r\n"
			"\t-p [number]: set balance port(default 80)\r\n"
			"\r\n"
			"version: v1.0\r\n");
}
/*
 * 包回调函数
 */
int cb(struct nfq_q_handle* qh, struct nfgenmsg* nfmsg,
		struct nfq_data* nfa, void* data)
{
	unsigned int					id			= 0;
	struct nfqnl_msg_packet_hdr*	ph			= NULL;
	unsigned char*					pdata		= NULL;
	int								pdata_len	= 0;
	struct iphdr*					piphdr		= NULL;
	struct tcphdr*					ptcphdr		= NULL;
	struct nfqnl_msg_packet_hw*		phw			= NULL;

	ph = nfq_get_msg_packet_hdr(nfa);
	if( ph )
	{
		id = ntohl(ph->packet_id);
	}

	pdata_len = nfq_get_payload(nfa, (char **)&pdata);
	if( -1 == pdata_len )
	{
		pdata_len = 0;
	}

	piphdr = (struct iphdr*)pdata;
	if( NULL == piphdr )
	{
		syslog(LOG_INFO, "nfq_get_payload failed: %s", strerror(errno));
		goto cb_end;
	}

	if( IPPROTO_TCP != piphdr->protocol )
	{
		goto cb_end;
	}

	if(0 != g_port)
	{
		ptcphdr = (struct tcphdr*)((unsigned char*)piphdr + (piphdr->ihl << 2));
		if(g_port != ptcphdr->dest)
		{
			goto cb_end;
		}
	}

	// 增加mark
	g_index = g_index % g_all_count + 1;
	return nfq_set_verdict_mark(qh, id, NF_REPEAT, htonl(g_index), pdata_len, pdata);
cb_end:
	return nfq_set_verdict_mark(qh, id, NF_ACCEPT, 0, pdata_len, pdata);
}

/*
 * 主函数入口
 */
void nfq_listen(void)
{
	struct nfq_handle*			h;
	struct nfq_q_handle*		qh;
	int							fd;
	int							rv;
	char						buf[40960];

	// 打开监控驱动
	h = nfq_open();
	if( !h )
	{
		syslog(LOG_INFO, "nfq_open failed: %s", strerror(errno));
		return;
	}
	// 注销绑定协议
	if( nfq_unbind_pf(h, AF_INET) < 0 )
	{
		syslog(LOG_INFO, "nfq_unbind_pf failed: %s", strerror(errno));
		return;
	}
	// 重新绑定协议
	if( nfq_bind_pf(h, AF_INET) < 0 )
	{
		syslog(LOG_INFO, "nfq_bind_pf failed: %s", strerror(errno));
		return;
	}
	// 创建监听队列
	qh = nfq_create_queue(h, 8800, &cb, NULL);
	if( !qh )
	{
		syslog(LOG_INFO, "nfq_create_queue failed: %s", strerror(errno));
		return;
	}

	if( nfq_set_mode(qh, NFQNL_COPY_PACKET, 0xffff) < 0 )
	{
		syslog(LOG_INFO, "nfq_set_mode failed: %s", strerror(errno));
		return;
	}
	// 获取文件描述符
	fd = nfq_fd(h);
	// 开始循环
	while( (rv = recv(fd, buf, sizeof(buf), 0)) >= 0 )
	{
		nfq_handle_packet(h, buf, rv);
	}
	// 退出
	syslog(LOG_INFO, "exit loop: %s", strerror(errno));
	nfq_destroy_queue(qh);
	nfq_close(h);
}


